#include "m_memory.h"

#include "esp_log.h"

static const char * TAG = "MEM";

// 初始化内存池
void mempool_init(MemoryPoolManager* mgr, void* buffer, size_t buffer_size)
{
    if (!mgr || !buffer || buffer_size < sizeof(MemoryBlock)) {
        ESP_LOGE("MEM", "Invalid input to mempool_init");
        return;
    }

    memset(mgr, 0, sizeof(MemoryPoolManager));
    mgr->pool_start = buffer;
    mgr->pool_size = buffer_size;
    mgr->free_call_count = 0;  // 初始化 free 调用计数器为 0

    MemoryBlock* block = (MemoryBlock*)buffer;
    block->canary_start = CANARY_VALUE;
    block->header_magic = BLOCK_HEADER_MAGIC;
    block->size = buffer_size - sizeof(MemoryBlock);
    block->is_allocated = 0;
    block->prev = NULL;
    block->next = NULL;
    block->footer_magic = BLOCK_FOOTER_MAGIC;
    block->canary_end = CANARY_VALUE;

    mgr->head = block;
    mgr->free_blocks = 1;
}

int is_valid_block(MemoryBlock* block)
{
    if (!block) return 0;
    if (block->canary_start != CANARY_VALUE || block->canary_end != CANARY_VALUE) {
        ESP_LOGE("MEM", "Canary corruption detected in block %p", block);
        return 0;
    }
    if (block->header_magic != BLOCK_HEADER_MAGIC ||
        block->footer_magic != BLOCK_FOOTER_MAGIC) {
        ESP_LOGE("MEM", "Header/footer magic corrupted in block %p", block);
        return 0;
    }
    return 1;
}

void merge_free_blocks(MemoryPoolManager* mgr, MemoryBlock* block)
{
    if (!block || !is_valid_block(block))
        return;

    if (block->next && block->next->is_allocated == 0 && is_valid_block(block->next)) {
        block->size += sizeof(MemoryBlock) + block->next->size;
        MemoryBlock* tmp = block->next;
        block->next = tmp->next;
        if (tmp->next) {
            tmp->next->prev = block;
        }
        mgr->free_blocks--;
    }

    if (block->prev && block->prev->is_allocated == 0 && is_valid_block(block->prev)) {
        block->prev->size += sizeof(MemoryBlock) + block->size;
        block->prev->next = block->next;
        if (block->next) {
            block->next->prev = block->prev;
        }
        mgr->free_blocks--;
    }
}

void* mempool_alloc(MemoryPoolManager* mgr, size_t size)
{
    if (!mgr || size == 0)
        return NULL;

    // 检查是否有足够的空间
    if (size > mgr->pool_size - mgr->total_alloc_size - sizeof(MemoryBlock)) {
        ESP_LOGE("MEM", "Not enough memory in pool for allocation of size: %zu", size);
        return NULL;
    }

    MemoryBlock* current = mgr->head;
    while (current != NULL) {
        if (!current->is_allocated && current->size >= size) {
            // 设置为已分配
            current->is_allocated = 1;

            // 如果剩余空间足够创建一个新的块
            if (current->size > size + sizeof(MemoryBlock)) {
                MemoryBlock* new_block = (MemoryBlock*)((char*)current + sizeof(MemoryBlock) + size);

                new_block->canary_start = CANARY_VALUE;
                new_block->header_magic = BLOCK_HEADER_MAGIC;
                new_block->size = current->size - size - sizeof(MemoryBlock);
                new_block->is_allocated = 0;
                new_block->prev = current;
                new_block->next = current->next;
                new_block->footer_magic = BLOCK_FOOTER_MAGIC;
                new_block->canary_end = CANARY_VALUE;

                if (new_block->next)
                    new_block->next->prev = new_block;

                current->size = size;
                current->next = new_block;

                mgr->free_blocks++;
            }

            current->footer_magic = BLOCK_FOOTER_MAGIC;
            current->canary_end = CANARY_VALUE;

            mgr->used_blocks++;
            mgr->total_alloc_size += size;

            return (void*)(current + 1); // 返回用户可用内存区域
        }
        current = current->next;
    }

    ESP_LOGE("MEM", "Failed to allocate memory of size: %zu", size);
    return NULL;
}

void mempool_free(MemoryPoolManager* mgr, void* ptr)
{
    if (!ptr || !mgr)
        return;

    MemoryBlock* block = (MemoryBlock*)ptr - 1;

    if (!is_valid_block(block)) {
        ESP_LOGE("MEM", "Attempted to free invalid or corrupted block at %p", block);
        return;
    }

    if (!block->is_allocated) {
        ESP_LOGE("MEM", "Double free detected at %p", block);
        return;
    }

    block->is_allocated = 0;
    mgr->used_blocks--;
    mgr->total_alloc_size -= block->size;

    merge_free_blocks(mgr, block);
    // 增加 free 调用计数
    mgr->free_call_count++;
    // 每调用 200 次 mempool_free 打印一次内存状态
    if (mgr->free_call_count % 200 == 0) {
        mempool_print_status(mgr);
    }
}

void mempool_print_status(MemoryPoolManager* mgr)
{
    if (!mgr) {
        ESP_LOGE("MEM", "Memory pool manager is NULL");
        return;
    }

    size_t used_bytes = 0;
    size_t free_bytes = 0;
    size_t allocated_count = 0;
    size_t free_count = 0;

    MemoryBlock* current = mgr->head;
    while (current != NULL) {
        if (is_valid_block(current)) {  // 确保当前块未损坏
            if (current->is_allocated) {
                used_bytes += current->size + sizeof(MemoryBlock);
                allocated_count++;
            } else {
                free_bytes += current->size + sizeof(MemoryBlock);
                free_count++;
            }
        }
        current = current->next;
    }

    ESP_LOGI("MEM", "=== Memory Pool Status ===");
    ESP_LOGI("MEM", "Total Pool Size:     %u bytes", (unsigned int)mgr->pool_size);
    ESP_LOGI("MEM", "Used Blocks:         %u blocks (%u bytes)",  (unsigned int)allocated_count, (unsigned int)used_bytes);
    ESP_LOGI("MEM", "Free Blocks:         %u blocks (%u bytes)",  (unsigned int)free_count, (unsigned int)free_bytes);
    ESP_LOGI("MEM", "Allocated Data Size: %zu bytes",  (unsigned int)mgr->total_alloc_size);
    ESP_LOGI("MEM", "=========================");
}